#-*- coding:utf-8 -*-
"""
Simple DB wrapper.
"""
from contextlib import contextmanager

import mysql.connector
from mysql.connector import errorcode

db_downline_config = {
    "host": 'cq01-dr-peercenterbak01.cq01',
    "user": "root",
    "password": "test123",
    "port": 5145,
    "autocommit": True,
    "buffered": True,
    "database": 'weigou_entity'
}

db_online_config = {
    "host": "10.52.120.76",
    "user": "wg_en",
    "password": "wg_en",
    "port": 7020,
    "autocommit": True,
    "buffered": True,
    "database": "weigou_entity"
}


def _get_connection(config):
    cnx = mysql.connector.connect(**config)
    return cnx


def get_connection(db_type, db_name=None):
    if db_name is not None:
        db_online_config["database"] = db_name
        db_downline_config["database"] = db_name
    else:
        db_online_config["database"] = "weigou_entity"
        db_downline_config["database"] = "weigou_entity"
    if db_type == "online":
        return _get_connection(db_online_config)
    elif db_type == "downline":
        return _get_connection(db_downline_config)
    else:
        raise Exception("no such db type:%s" % db_type)


@contextmanager
def get_db_cursor(cur_type, db_name=None):
    assert cur_type in ("online", "downline")
    cnx = get_connection(cur_type, db_name)
    cursor = cnx.cursor()
    try:
        yield cursor
    except:
        cnx.rollback()
        raise
    else:
        cnx.commit()
    finally:
        cursor.close()
        cnx.close()


if __name__ == "__main__":
    with get_db_cursor("downline") as cursor:
        sql = "select * from test"
        cursor.execute(sql)
        for row in cursor:
            print row